package security;

import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

import org.bouncycastle.util.encoders.Base64;

/**
 * Provides RSA-encrypted communication.
 * 
 * @author Lukas
 *
 */
public class RSAChannel extends ChannelDecorator {
	private Cipher encrypt, decrypt;
	
	/**
	 * Constructor.
	 * @param channel The underlying channel that shall be used for transmission.
	 * @param pubKey The public Key (for encryption).
	 * @param priKey The private Key (for decryption). If <tt>null</tt>, the private key will be loaded on demand.
	 * @throws IOException
	 */
	public RSAChannel(Channel channel, PublicKey pubKey, PrivateKey priKey) throws IOException {
		this.channel = channel;
		initCiphers(pubKey, priKey);
	}

	@Override
	public void send(String message) throws IOException {
		try {
			channel.send(new String(Base64.encode(encrypt.doFinal(message.getBytes()))));
		} catch (IllegalBlockSizeException e) {
			throw new IOException(e);
		} catch (BadPaddingException e) {
			throw new IOException(e);
		}
	}

	@Override
	public String receive() throws IOException {
		String message = channel.receive();
		if (message.startsWith("!list")) {
			return message;
		}
		String decr;
		try {
			decr = new String(decrypt.doFinal(Base64.decode(message)));
		} catch (IllegalBlockSizeException e) {
			return message;
		} catch (BadPaddingException e) {
			return message;
		}
		if (decr.startsWith("!login ") || decr.startsWith("!ok ")) {
			return decr;
		} else {
			return message;
		}
	}

	private void initCiphers(PublicKey pubKey, PrivateKey priKey) throws IOException {
		try {
			if (pubKey != null)
				encrypt = Cipher.getInstance("RSA/NONE/OAEPWithSHA256AndMGF1Padding");
			if (priKey != null)
				decrypt = Cipher.getInstance("RSA/NONE/OAEPWithSHA256AndMGF1Padding");
		} catch (NoSuchAlgorithmException e) {
			throw new IOException(e);
		} catch (NoSuchPaddingException e) {
			throw new IOException(e);
		}
		try {
			if (pubKey != null)
				encrypt.init(Cipher.ENCRYPT_MODE, pubKey);
			if (priKey != null)
				decrypt.init(Cipher.DECRYPT_MODE, priKey);
		} catch (InvalidKeyException e) {
			throw new IOException(e);
		}
	}
}
